from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
from langchain_huggingface import HuggingFacePipeline
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import torch

# 加载大模型
model_name = "/data/ai/LLaMA-Factory/model/Qwen2.5-7B-Instruct/qwen/Qwen2___5-7B-Instruct/"
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16,
    device_map='cuda:0'  # 使用GPU进行推理
)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# 设置生成管道
gen_pipeline = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    max_length=1024,
    do_sample=True,
    top_k=10,
    truncation=True
)

# 使用 langchain 包装生成器
llm = HuggingFacePipeline(pipeline=gen_pipeline)

# FastAPI 初始化
# app = FastAPI()

# 定义请求模型
class PromptRequest(BaseModel):
    prompt: str
    system_prompt: str  # 添加系统提示词字段



full_prompt = f"系统提示词: {'输出限制在100token以内'}\n用户问题: {'你是一个热心且专业的客服代表，专注于回答与公司产品、服务和技术支持相关的问题。'}"

# 通过 langchain 获取生成的答案
output = llm.invoke(full_prompt)
print(output)  # 返回生成的文本
# 创建API路由，接收用户输入的 prompt 和系统提示词，返回模型生成的结果
# @app.post("/generate_with_system_prompt")
# async def generate_text_with_system_prompt(request: PromptRequest):
#     try:
        # 拼接系统提示词和用户输入的prompt
        # full_prompt = f"系统提示词: {request.system_prompt}\n用户问题: {request.prompt}"
        #
        # # 通过 langchain 获取生成的答案
        # output = llm.invoke(full_prompt)
        # return {"generated_text": output[0]['generated_text']}  # 返回生成的文本
    # except Exception as e:
    #     raise HTTPException(status_code=500, detail=str(e))

# 运行命令: uvicorn test1:app --host 0.0.0.0 --port 50001
